# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
# Set up custom environment before nearly anything else is imported
# NOTE: this should be the first import (no not reorder)

import argparse
import os

import torch
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.data import make_data_loader
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer
from maskrcnn_benchmark.utils.collect_env import collect_env_info
from maskrcnn_benchmark.utils.comm import synchronize, get_rank
from maskrcnn_benchmark.utils.logger import setup_logger
from maskrcnn_benchmark.utils.miscellaneous import mkdir

from maskrcnn_benchmark.modeling.box_coder import BoxCoder
from maskrcnn_benchmark.modeling.rpn.anchor_generator import make_anchor_generator_retinanet
from maskrcnn_benchmark.modeling.rpn.retinanet import RetinaNetHead
from nms_recall.inference_for_NR import inference
from nms_recall.without_nms_postprocessor import make_retinanet_postprocessor


class RetinaNetModule(torch.nn.Module):
	"""
	Module for RetinaNet computation. Takes feature maps from the backbone and RPN
	proposals and losses.
	"""

	def __init__(self, cfg):
		super(RetinaNetModule, self).__init__()

		self.cfg = cfg.clone()

		anchor_generator = make_anchor_generator_retinanet(cfg)
		head = RetinaNetHead(cfg)
		box_coder = BoxCoder(weights=(10., 10., 5., 5.))

		if self.cfg.MODEL.SPARSE_MASK_ON:
			raise NotImplementedError
		else:
			box_selector_test = make_retinanet_postprocessor(
				cfg, 100, box_coder)

		self.anchor_generator = anchor_generator
		self.head = head
		self.box_selector_test = box_selector_test

	def forward(self, images, features, target=None):
		box_cls, box_regression = self.head(features)
		anchors = self.anchor_generator(images, features)

		if self.training:
			raise NotImplementedError
		else:
			return self._forward_test(anchors, box_cls, box_regression)

	def _forward_test(self, anchors, box_cls, box_regression):
		boxes = self.box_selector_test(anchors, box_cls, box_regression)
		return (anchors, boxes), {}


def main():
	parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference")
	parser.add_argument(
		"--config-file",
		default="configs/free_anchor_R-50-FPN_8gpu_1x.yaml",
		metavar="FILE",
		help="path to config file",
	)
	parser.add_argument("--local_rank", type=int, default=0)
	parser.add_argument(
		"opts",
		help="Modify config options using the command-line",
		default=None,
		nargs=argparse.REMAINDER,
	)

	args = parser.parse_args()

	num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
	distributed = num_gpus > 1

	if distributed:
		torch.cuda.set_device(args.local_rank)
		torch.distributed.deprecated.init_process_group(
			backend="nccl", init_method="env://"
		)

	cfg.merge_from_file(args.config_file)
	cfg.merge_from_list(args.opts)
	# cfg.merge_from_list(['TEST.IMS_PER_BATCH', 1])
	cfg.freeze()

	save_dir = ""
	logger = setup_logger("maskrcnn_benchmark", save_dir, get_rank())
	logger.info("Using {} GPUs".format(num_gpus))
	logger.info(cfg)

	logger.info("Collecting env info (might take some time)")
	logger.info("\n" + collect_env_info())

	model = build_detection_model(cfg)

	model.rpn = RetinaNetModule(cfg)

	model.to(cfg.MODEL.DEVICE)

	checkpointer = DetectronCheckpointer(cfg, model)
	_ = checkpointer.load(cfg.MODEL.WEIGHT)

	iou_types = ("bbox",)
	output_folders = [None] * len(cfg.DATASETS.TEST)
	if cfg.OUTPUT_DIR:
		dataset_names = cfg.DATASETS.TEST
		for idx, dataset_name in enumerate(dataset_names):
			output_folder = os.path.join(cfg.OUTPUT_DIR, "NR", dataset_name)
			mkdir(output_folder)
			output_folders[idx] = output_folder
	data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed)
	for output_folder, data_loader_val in zip(output_folders, data_loaders_val):
		inference(
			model,
			data_loader_val,
			iou_types=iou_types,
			# box_only=cfg.MODEL.RPN_ONLY,
			box_only=False if cfg.RETINANET.RETINANET_ON else cfg.MODEL.RPN_ONLY,
			device=cfg.MODEL.DEVICE,
			expected_results=cfg.TEST.EXPECTED_RESULTS,
			expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL,
			output_folder=output_folder,
		)
		synchronize()


if __name__ == "__main__":
	main()
